from models.resnet import resnet18
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import os
from utils.parse_config import parse_data_config
from data.BottleDataset import BottleDataset
import argparse
try:
    import moxing as mox
except:
    print('not use moxing')


class Trainer(object):
    def __init__(self, learning_rate=0.001, n_class=13,device=torch.device('cpu')):
        self.model = resnet18(n_class)

        # self.optimizer = torch.optim.SGD(self.model.parameters(), lr=0.001, momentum=0.5,
        #                                  weight_decay=0, nesterov=True)
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=learning_rate)
        # self.optimizer = torch.optim.Adadelta(self.model.parameters(), lr=learning_rate, rho=0.9, eps=1e-06, weight_decay=0)
        self.criterion =nn.CrossEntropyLoss()
        self.n_class = n_class
        self.train_local=train_local
        self.device =device
        print("---optimizer:",self.optimizer)
        print("---criterion:",self.criterion)
        # viz.line([0.], [0], win='train_loss', opts=dict(title='train_loss'))
        # viz.line([0.], [0], win='miou', opts=dict(title='miou'))
        # self.viz = viz

    def train(self, epoch, train_loader, n_data):
        train_loss = 0.0
        correct_count = torch.tensor(0,dtype=torch.int8)
        for i, training_data in enumerate(train_loader, 0):
            inputs, labels = training_data[0], training_data[1]
            inputs = inputs.to(self.device)
            labels = labels.to(self.device)
            self.model.to(self.device)
            self.optimizer.zero_grad()
            output= self.model(inputs)
            loss = self.criterion(output, labels)


            loss.backward()
            self.optimizer.step()
            train_loss += loss.item()

            pred = torch.argmax(output.data,dim=1)
            result = labels == pred
            correct_count+=torch.sum(result).int()
        print('---[Epoch: %d]' % (epoch))
        print('---Training Loss: %.3f' % train_loss)
        print("---Training Accuracy: %i/%i=%f"%(correct_count,n_data,torch.true_divide(correct_count,n_data)))

    def validation(self, epoch, validat_loader, n_data):
        validation_loss = 0.0
        correct_count = 0
        for i, training_data in enumerate(validat_loader, 0):
            inputs, labels = training_data[0], training_data[1]
            inputs = inputs.to(self.device)
            labels = labels.to(self.device)
            self.model.to(self.device)
            with torch.no_grad():
                output = self.model(inputs)
            loss = self.criterion(output, labels)
            validation_loss += loss.item()
            pred = torch.argmax(output.data,dim=1)
            result = labels == pred
            correct_count+=torch.sum(result).int()
        print('---[Epoch: %d]' % (epoch))
        print('---Validation Loss: %.3f' % validation_loss)
        print("---Validation Accuracy: %i/%i=%f"%(correct_count,n_data,torch.true_divide(correct_count,n_data)))

    def test(self,  test_loader, n_data):
        test_loss = 0.0
        correct_count = 0
        for i, training_data in enumerate(test_loader, 0):
            inputs, labels = training_data[0], training_data[1]
            inputs = inputs.to(self.device)
            labels = labels.to(self.device)
            self.model.to(self.device)
            with torch.no_grad():
                output = self.model(inputs)
            loss = self.criterion(output, labels)
            test_loss += loss.item()
            pred = torch.argmax(output.data,dim=1)
            result = labels == pred
            correct_count+=torch.sum(result).int()

        print('---Testing Loss: %.3f' % test_loss)
        print("---Testing Accuracy: %i/%i=%.4f"%(correct_count,n_data,torch.true_divide(correct_count,n_data)))

def prepare_data_on_modelarts(pretrained_weights, local_data_root, data_url, train_url):
    """
    将OBS上的数据拷贝到ModelArts中
    """
    # 拷贝预训练参数文件

    # 默认使用ModelArts中的如下两个路径用于存储数据：
    # 0) /cache/model: 如果适用预训练模型，存储从OBS拷贝过来的预训练模型
    # 1）/cache/datasets: 存储从OBS拷贝过来的训练数据
    # 2）/cache/log: 存储训练日志和训练模型，并且在训练结束后，该目录下的内容会被全部拷贝到OBS
    if pretrained_weights:
        _, weights_name = os.path.split(pretrained_weights)
        mox.file.copy(pretrained_weights, os.path.join(local_data_root, 'model/' + weights_name))
        pretrained_weights = os.path.join(local_data_root, 'model/' + weights_name)

    data_local = os.path.join(local_data_root, 'datasets/BottleDataset')
    if not os.path.exists(data_local):
        data_dir = os.path.join(local_data_root, 'datasets')
        print("data_dir:",data_dir)
        mox.file.copy_parallel(data_url, data_dir)
        os.system('cd %s;unzip BottleDataset.zip' % data_dir)  # 训练集已提前打包为BottleDataset.zip
        print("data_local:",data_local)
        if os.path.isdir(data_local):
            os.system('cd %s;rm BottleDataset.zip' % data_dir)
            print('unzip BottleDataset.zip success, args.data_local is', data_local)
        else:
            raise Exception('unzip BottleDataset.zip Failed')
    else:
        print('args.data_local: %s is already exist, skip copy' % data_local)

    output_path = os.path.join(local_data_root, 'log/')
    if not os.path.exists(output_path):
        os.mkdir(output_path)

    return pretrained_weights, data_local, output_path


if __name__ == "__main__":
    # python -m visdom.server
    parser = argparse.ArgumentParser()

    parser.add_argument('--data_url', default="", type=str, help='the training and validation data path')
    parser.add_argument('--train_url', default="", type=str, help='the path to save training outputs')
    parser.add_argument("--pretrained_weights", type=str,default="", help="if specified starts from checkpoint model")

    parser.add_argument("--epoch", type=int,default=105, help="learning_rate")
    parser.add_argument("--batch_size", type=int,default=64, help="batch size")
    parser.add_argument("--learning_rate", type=float,default=0.0005, help="learning_rate")
    opt = parser.parse_args()
    print(opt)
    train_local = opt.data_url==''
    EPOCH = opt.epoch
    BATCH_SIZE = 4 if train_local else opt.batch_size
    learning_rate = opt.learning_rate
    base_size = 256
    print("---EPOCH",EPOCH)
    print("---BATCH_SIZE",BATCH_SIZE)
    print("---learning_rate",learning_rate)
    print("---base_size",base_size)

    # train_local = True
    if train_local:
        # data_path = '../RoadLineDataset'
        data_config_path = 'data/custom.data'
        pretrained_weights = ''
        data_local='F:/GitRepository/BottleClassification/BottleDataset/'
        output_path=''
    else:
        # Get data configuration
        data_config_path = 'LaneLineDetection/data/custom.data'
        local_data_root = '/cache/'
        # data_url = 'obs://class-1275-42261/Lab-2210/modelarts22905308/datasets/'
        data_url=opt.data_url
        # train_url = 'obs:///class-1275-42261/Lab-2210/modelarts22905308/train_output/V0015/'
        train_url=opt.train_url
        pretrained_weights = opt.pretrained_weights
        pretrained_weights, data_local, output_path = prepare_data_on_modelarts(pretrained_weights, local_data_root,
                                                                                data_url, train_url)
        print("output_path:",output_path)
        print("data_local:", data_local)
    data_config = parse_data_config(data_config_path)
    train_path = data_config["local_train"] if train_local else data_config["train"]
    # valid_path = data_config["local_valid"] if train_local else data_config["valid"]
    train_data = BottleDataset(train_path,data_local, base_size=base_size)
    print("---Dataset:",train_data.__len__())
    # valid_data = BottleDataset(valid_path, base_size=base_size)
    train_data_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)  # 使用DataLoader加载数据
    # valid_data_loader = DataLoader(valid_data, batch_size=BATCH_SIZE, shuffle=True)  # 使用DataLoader加载数据
    # tbar = tqdm(data_loader)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("device:", device)
    trainer = Trainer(learning_rate=learning_rate, device=device)

    for epoch in range(1,EPOCH+1):
        trainer.train(epoch, train_data_loader, train_data.__len__())
        # trainer.validation(epoch, data_loader, BATCH_SIZE, device)
        if epoch % 5 == 0:
            # AP = trainer.validation(epoch, valid_data_loader, BATCH_SIZE, device)
            temp_model_name = f"ckpt_%d.pth" % (epoch)
            ckpt_name = os.path.join('models', 'save', temp_model_name)
            if not train_local:
                ckpt_name = os.path.join('LaneLineDetection', ckpt_name)
            torch.save(trainer.model.state_dict(), ckpt_name)
            # save model to obs
            if not train_local:
                mox.file.copy_parallel(ckpt_name, os.path.join(train_url, temp_model_name))

            # Test
            test_path = data_config["local_test"] if train_local else data_config["test"]

            test_data = BottleDataset(test_path,data_local, base_size=base_size,is_train=False)

            test_data_loader = DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=True)  # 使用DataLoader加载数据

            trainer.test(test_data_loader,test_data.__len__())

